
package com.universest.swordholder.aitool

import android.util.Log
import com.alibaba.fastjson.JSON
import com.alibaba.fastjson.JSONArray
import com.alibaba.fastjson.JSONObject
import com.universest.swordholder.network.aiHttpClient
import okhttp3.*
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.RequestBody.Companion.toRequestBody
import okhttp3.sse.EventSource
import okhttp3.sse.EventSourceListener
import okhttp3.sse.EventSources
import java.io.IOException
import java.util.concurrent.TimeUnit

class DeepSeekApi(private val apiKey: String,private val deepseekModel: String,val url: String = Constants.DEEPSEEK_API_URL) : AbstractLlmApi() {

    @Throws(IOException::class)
    override fun chat(messages: List<Message>): Message? {
        val requestJson = buildRequestBody(messages, false)
        val request = buildHttpRequest(requestJson)
        aiHttpClient.newCall(request).execute().use { response ->
            if (!response.isSuccessful || response.body == null) {
                throw IOException("HTTP request failed with code: ${response.code}")
            }

            val responseBody = response.body!!.string()
            val jsonResponse = JSONObject.parseObject(responseBody)

            if (jsonResponse == null) {
                throw IOException("Failed to parse response JSON.")
            }

            val choices = jsonResponse.getJSONArray("choices")
            if (choices != null && choices.isNotEmpty()) {
                val choice = choices.getJSONObject(0)
                if (choice.containsKey("message")) {
                    val messageJson = choice.getJSONObject("message")
                    return parseJsonToMessage(messageJson)
                }
            }
        }
        return null
    }

    private fun buildHttpRequest(requestJson: JSONObject): Request {
        val jsonString = requestJson.toString()
        return Request.Builder()
            .url("$url/chat/completions")
            .addHeader("Content-Type", "application/json")
            .addHeader("Authorization", "Bearer $apiKey")
            .post(jsonString.toRequestBody("application/json".toMediaType()))
            .build()
    }

    private fun parseJsonToMessage(jsonObject: JSONObject): Message {
        val role = jsonObject.getOrDefault("role", "user").toString()
        val content = jsonObject.getOrDefault("content", "").toString()
        val deepThought = jsonObject.getOrDefault("reasoning_content", "").toString()
        return Message(role, content, deepThought)
    }

    private fun buildRequestBody(messages: List<Message>, stream: Boolean): JSONObject {
        val jsonObject = JSONObject()
        jsonObject["model"] = getModel()
        jsonObject["messages"] = wrapMessages(messages)
        jsonObject["max_tokens"] = maxTokens
        jsonObject["temperature"] = temperature
        jsonObject["top_p"] = topP
        jsonObject["presence_penalty"] = presencePenalty
        jsonObject["stream"] = stream
        jsonObject["frequency_penalty"] = frequencyPenalty
        jsonObject["response_format"] = JSON.parse(responseFormat)

        if (stops != null && stops!!.isNotEmpty()) {
            if (stops!!.size == 1) {
                jsonObject["stop"] = stops!![0]
            } else {
                jsonObject["stop"] = stops
            }
        }

        jsonObject["stream_options"] = null
        jsonObject["tools"] = null
        jsonObject["tool_choice"] = "none"
        jsonObject["logprobs"] = false
        jsonObject["top_logprobs"] = null

        return jsonObject
    }

    override fun chatStream(messages: List<Message>, callback: StreamCallback): EventSource {
        val requestBody = buildRequestBody(messages,true)
        val request = buildHttpRequest(requestBody)
        val factory = EventSources.createFactory(aiHttpClient)
        return factory.newEventSource(request, object : EventSourceListener() {
            override fun onClosed(eventSource: EventSource) {
                super.onClosed(eventSource)
                callback.onComplete()
            }

            override fun onEvent(
                eventSource: EventSource,
                id: String?,
                type: String?,
                data: String
            ) {
                super.onEvent(eventSource, id, type, data)
                try {
                    Log.d("DeepSeekApi", "onEvent: $data")
                    if(data == "[DONE]"){
                        return
                    }
                    val jsonObject = JSON.parseObject(data)
                    val choices = jsonObject.getJSONArray("choices")
                    if (choices != null && choices.isNotEmpty()) {
                        for(i in 0 until choices.size){
                            val choice = choices.getJSONObject(i)
                            val delta = choice.getJSONObject("delta")
                            if(delta != null) {
                                val content = delta.getOrDefault("content", "")?.toString()?:""
                                if (content.isNotEmpty()) {
                                    callback.onChunk(content)
                                }
                                val reasoning =
                                    delta.getOrDefault("reasoning_content", "")?.toString()?:""
                                if (reasoning.isNotEmpty()) {
                                    callback.onDeepThought(reasoning)
                                }
                            }
                        }
                    }
                }catch (e: Exception){
                    Log.e("DeepSeekApi", "onEvent error: $e")
                }
            }

            override fun onFailure(
                eventSource: EventSource,
                t: Throwable?,
                response: Response?
            ) {
                super.onFailure(eventSource, t, response)
                callback.onError(if(t != null) Exception(t) else Exception(response?.message?:"Unknown error"))
            }

            override fun onOpen(eventSource: EventSource, response: Response) {
                super.onOpen(eventSource, response)
            }
        })
    }

    override fun getModel(): String = deepseekModel

    private fun wrapMessages(messages: List<Message>): JSONArray {
        val jsonArray = JSONArray()
        for (message in messages) {
            val jsonObject = JSONObject()
            jsonObject["role"] = message.role
            jsonObject["content"] = message.content
            jsonArray.add(jsonObject)
        }
        return jsonArray
    }
}